module CSRMatrix_m
use SMatrix_m
type CSRMatrix
    integer :: m, n
    logical :: isSym
    integer :: NNZ
    integer, allocatable :: rowIndex(:), columns(:)
    real*8, allocatable :: values(:)
contains
    procedure :: initCSRMatrix
    procedure :: mulVector
    procedure :: mulVectorS
    procedure :: deallocateCSRMatrix
end type
interface operator(*)
    module procedure :: CSRmul
end interface

contains

subroutine initCSRMatrix(this, A)
    implicit none
    class(CSRMatrix), intent(out) :: this
    type(SMatrix), intent(inout) :: A
    
    this.m = A%M
    this.n = A%N
    if ( .NOT. A.isSorted) call A.sort()
    this.isSym = A.isSymmetric()
    call A.exportCSRFormat(this.isSym, this.rowIndex, this.columns, this.values)
    this.NNZ = size(this.columns, 1)
end subroutine

subroutine SMatrixToCSRMatrix(A, diag, mat1, mat2)
    implicit none
    type(SMatrix), intent(inout) :: A
    real*8, intent(inout) :: diag(:)
    type(CSRMatrix), intent(out) :: mat1, mat2
    
    integer :: m, n, nnzL, nnzU, nnz1, nnz2
    integer :: rowIdxU, rowIdxL, colIdxL, colIdxU
    integer :: i, j
    type(SEntry), pointer :: p
    m = A%M
    n = A%N
    diag = 0.0D0
    mat1%M = m; mat1%N = n
    mat2%M = m; mat2%N = n
    if (.NOT. A.isSorted) call A.sort()
    call A.getNNZLU(nnzL, nnzU)
    allocate(mat1.rowIndex(m+1), mat2.rowIndex(m+1))
    mat1.rowIndex = 0
    mat2.rowIndex = 0
    if (A.isSymmetric()) then
        nnz1 = m + nnzU
        nnz2 = nnzL
        mat1.NNZ = nnz1; mat2.NNZ = nnz2
        allocate(mat1.columns(nnz1), mat1.values(nnz1))
        allocate(mat2.columns(nnz2), mat2.values(nnz2))
        colIdxU = 0; colIdxL = 0
        mat1.rowIndex(1) = 1
        mat2.rowIndex(1) = 1
        do i = 1, M
            rowIdxU = 0
            rowIdxL = 0
            p => A%row(i)%head
            do while(associated(p))
                if (p%col < i) then
                    colIdxL = colIdxL + 1
                    rowIdxL = rowIdxL + 1
                    mat2.columns(colIdxL) = p%col
                    mat2.values(colIdxL) = p%value
                else
                    colIdxU = colIdxU + 1
                    rowIdxU = rowIdxU + 1
                    mat1.columns(colIdxU) = p%col
                    mat1.values(colIdxU) = p%value                    
                end if
                if (p%col == i) then
                    diag(i) = p%value
                end if
                p => p%next
            end do
            mat1.rowIndex(i + 1) = mat1.rowIndex(i) + rowIdxU; 
 			mat2.rowIndex(i + 1) = mat2.rowIndex(i) + rowIdxL;          
        end do
        mat1.isSym = .TRUE.
    else
        nnz1 = A%NNZ; nnz2 = nnzL
        mat1.NNZ = nnz1; mat2.NNZ = nnz2
        allocate(mat1.columns(nnz1), mat1.values(nnz1))
        allocate(mat2.columns(nnz2), mat2.values(nnz2))
        colIdxU = 0; colIdxL = 0
        mat1.rowIndex(1) = 1
        mat2.rowIndex(1) = 1
        do i = 1, M
            rowIdxU = 0
            rowIdxL = 0
            p => A%row(i)%head
            do while(associated(p))
                if (p%col < i) then
                    colIdxL = colIdxL + 1
                    rowIdxL = rowIdxL + 1
                    mat2.columns(colIdxL) = p%col
                    mat2.values(colIdxL) = p%value
                end if
                colIdxU = colIdxU + 1
                rowIdxU = rowIdxU + 1
                mat1.columns(colIdxU) = p%col
                mat1.values(colIdxU) = p%value
                if (p%col == i) then
                    diag(i) = p%value
                end if
                p => p%next
            end do
 			mat2.rowIndex(i + 1) = mat2.rowIndex(i) + rowIdxL;
			mat1.rowIndex(i + 1) = mat1.rowIndex(i) + rowIdxU;           
        end do
        mat1.isSym = .FALSE.
    end if
end subroutine

function CSRmul(A, u) result(res)
    implicit none
    class(CSRMatrix), intent(in) :: A
    real*8, intent(in) :: u(A%N)
    real*8 :: res(A%M)
    if (A.isSym) then
        call mulVectorS(A, u, res)
    else
        call mulVector(A, u, res)
    end if
end function


subroutine mulVector(this, u, v)
    implicit none
    class(CSRMatrix), intent(in) :: this
    real*8, intent(in) :: u(:)
    real*8, intent(out) :: v(:)
    
    integer :: i, j, jStart, jEnd, colj
    do i = 1, this.m
        v(i) = 0.0D0
        jStart = this.rowIndex(i)
        jEnd = this.rowIndex(i + 1) - 1
        do j = jStart, jEnd
            v(i) = v(i) + this.values(j) * u(this.columns(j))
        end do
 !       v(i) = dot_product(this.values(jStart:jEnd), u(this.columns(jStart:jEnd)))
    end do
end subroutine

subroutine mulVectorS(this, u, v)
    implicit none
    class(CSRMatrix), intent(in) :: this
    real*8, intent(in) :: u(:)
    real*8, intent(out) :: v(:)
    
    integer :: i, j, jStart, jEnd, colj     
    v = 0.0D0
    do i = 1, this.m
        jStart = this.rowIndex(i)
        jEnd = this.rowIndex(i + 1) - 1
        do j = jStart, jEnd
            colj = this.columns(j)
            v(i) = v(i) +  this.values(j) * u(colj)
            if (i /= colj) then
                v(colj) = v(colj) +  this.values(j) * u(i)
            end if
        end do
        !v(i) = dot_product(this.values(jStart:jEnd), u(this.columns(jStart:jEnd)))
        !v(this.columns(jStart:jEnd)) = v(this.columns(jStart:jEnd)) + u(i) * this.values(jStart:jEnd)
        !v(i) = v(i) - this.
    end do
end subroutine

subroutine deallocateCSRMatrix(this)
    implicit none
    class(CSRMatrix), intent(out) :: this    
    this.isSym = .FALSE.
    this.m = 0
    this.n = 0
    this.NNZ = 0
    if (allocated(this.rowIndex)) deallocate(this.rowIndex)
    if (allocated(this.columns)) deallocate(this.columns)
    if (allocated(this.values)) deallocate(this.values)
end subroutine


end module
    